import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

import json
from typing import Optional, List, Dict

from datasets import load_dataset, DatasetDict


def human_to_user(examples):
    role_map = {'human': 'user', 'assistant': 'assistant'}
    return [{'role': role_map[example['role']], 'content': example['text']} for example in examples]

def row_map(example):
    return {
        'chosen': human_to_user(example['context'] + [example['chosen']]),
        'rejected': human_to_user(example['context'] + [example['rejected']]),
    }

root_path = r'C:\Users\james\.cache\huggingface\hub\datasets--dikw--hh_rlhf_cn\snapshots\82fdd2b7f454e62715e0d7d940cafcad081a66fb'

train_dataset = load_dataset('json', data_files=[os.path.join(root_path, 'hh_rlhf_train.jsonl')])
test_dataset = load_dataset('json', data_files=[os.path.join(root_path, 'hh_rlhf_test.jsonl')])

dataset = DatasetDict({
    'train': train_dataset['train'],
    'test': test_dataset['train'],
})

# dataset = load_dataset("dikw/hh_rlhf_cn")
dataset = dataset.map(row_map, remove_columns=['context'], desc="convert context to chosen and rejected")
print(dataset)

for split in ["train", "test"]:
    with open(f'hh_rlhf_cn_{split}.jsonl', 'w', encoding='utf-8') as f:
        for row in dataset[split]:
            f.write(json.dumps(row, ensure_ascii=False) + '\n')

print("end")